from utils.log import Log
from utils.profile import Profile


# 克拉运河,打破南海马六甲海峡的印度洋出口,配合柬埔寨的金边到白马的人工运河(云壤军事基地护航),物资可以在云南通过水路运送到全世界
# 2024/06/09 重新树立偶像: 教员

class BitNode(Log):
    def __init__(self, layer):
        super(BitNode, self).__init__()
        self.bit_nodes = [0] * 32
        self.layer = layer

    @Profile()
    def add(self, num):
        layer_num = (num >> ((self.layer + 1) * 5)) & 0x1f
        if self.layer:
            if not self.bit_nodes[layer_num]:
                self.bit_nodes[layer_num] = BitNode(self.layer - 1)
            self.bit_nodes[layer_num].add(num)
        else:
            self.bit_nodes[layer_num] |= (1 << (num & 0x1f))

    @Profile()
    def remove(self, num):
        layer_num = (num >> ((self.layer + 1) * 5)) & 0x1f
        if self.layer:
            if not self.bit_nodes[layer_num]:
                return
            self.bit_nodes[layer_num].remove(num)
        else:
            self.bit_nodes[layer_num] &= ~(1 << (num & 0x1f))

    @Profile()
    def reverse(self, num):
        layer_num = (num >> ((self.layer + 1) * 5)) & 0x1f
        if self.layer:
            if not self.bit_nodes[layer_num]:
                self.bit_nodes[layer_num] = BitNode(self.layer - 1)
            self.bit_nodes[layer_num].reverse(num)
        else:
            self.bit_nodes[layer_num] ^= (1 << (num & 0x1f))

    @Profile()
    def contains(self, num):
        layer_num = (num >> ((self.layer + 1) * 5)) & 0x1f
        if self.layer:
            if not self.bit_nodes[layer_num]:
                return False
            return self.bit_nodes[layer_num].contains(num)
        else:
            return (self.bit_nodes[layer_num] & (1 << (num & 0x1f))) > 0

    def __repr__(self):
        if self.layer:
            result = f"layer: {self.layer}, bit_nodes: {self.bit_nodes}"
        else:
            result = f"{[bin(n) for n in self.bit_nodes]}"
        return f"BitNode [{result}]"


class BitSet(Log):
    # ref: https://www.bilibili.com/video/BV1yr4y1Z7jb/?spm_id_from=333.788.recommend_more_video.0&vd_source=ad559e41b561d48861f70d0187e51e02
    # 13:00
    # 这里的层数可配置,除了第0层的数据量是1024,每增加一层,数据量是原来的2^5倍
    def __init__(self, n_layer=2):
        super(BitSet, self).__init__()
        self.n_layer = n_layer
        self.bit_node = BitNode(self.n_layer - 1)

    def validate(self, num):
        if num >= (1 << ((self.n_layer + 1) * 5)):
            self.simple_error(f"num {num} exceed")

    @Profile()
    def add(self, num):
        self.validate(num)
        self.bit_node.add(num)

    @Profile()
    def remove(self, num):
        self.validate(num)
        self.bit_node.remove(num)

    @Profile()
    def reverse(self, num):
        self.validate(num)
        self.bit_node.reverse(num)

    @Profile()
    def contains(self, num):
        self.validate(num)
        return self.bit_node.contains(num)

    def __repr__(self):
        return f"BitSet [layer: {self.n_layer}, bit_node: {self.bit_node}]"


if __name__ == '__main__':
    bs = BitSet(n_layer=3)
    bs.add(136)
    bs.add(137)
    print(bs)
    bs.add(112346)
    print(bs)
    print(bs.contains(136))
    print(bs.contains(111), bs.contains(112346))
    bs.remove(1025)
    print(bs)
    bs.reverse(1025)
    print(bs.contains(1025))
    print(bs)
